You write custom CUDA kernels to replace the PyTorch operators in the given EvoNorm architecture to get speedups.
You have complete freedom to choose the set of operators you want to replace. You may make the decision to replace some operators with custom CUDA kernels and leave others unchanged. You may replace multiple operators with custom implementations, consider operator fusion opportunities (combining multiple operators into a single kernel, for example, combining normalization+affine_transform+nonlinear_gating), or algorithmic changes (such as optimized memory access patterns). You are only limited by your imagination.

Circle Loss CUDA Optimization with Two-Stage Reduction

Core Algorithm
Loss Function: Implements Circle Loss for deep metric learning

Key Formula:

Positive pairs: logit_p = -α_p * (s_p - Δ_p) * γ

Negative pairs: logit_n = α_n * (s_n - Δ_n) * γ

Final loss: log(1 + exp(logsumexp_p + logsumexp_n))

CUDA Optimization Techniques
1. Two-Stage Parallel Reduction
Stage 1: Find global maximum of logits for numerical stability

Kernel: circleloss_find_max_kernel

Block-level max reduction with shared memory

Grid-stride loops for load balancing

Stage 2: Compute sum of exponentials with stability

Kernel: circleloss_sum_exp_diff_kernel

Subtract global max before exponentiation

Block-level sum reduction

2. Memory Access Optimization
Coalesced Memory Access: Sequential memory access patterns

Shared Memory Utilization:

s_data_p[BLOCK_SIZE] for positive pairs

s_data_n[BLOCK_SIZE] for negative pairs

Contiguous Tensors: Ensure input tensors are contiguous

3. Numerical Stability Features
Log-Sum-Exp Trick: Subtract maximum before exponentiation

Stable Softplus: max(0,x) + log(1 + exp(-abs(x)))

Float Safety: Use FLT_MAX for initialization

4. Parallelization Strategy
Grid-Stride Loops: Handle arbitrary input sizes

Block Size: 256 threads per block for optimal occupancy

Dynamic Grid Size: (n_elements + BLOCK_SIZE - 1) / BLOCK_SIZE

5. Kernel Design Patterns

6. Performance Optimizations
Minimal CPU-GPU Synchronization: Avoid unnecessary .item() calls

Batch Matrix Multiplication: Precompute similarity matrix

Memory Pre-allocation: Pre-allocate output buffers

7. Implementation Features
Template-Free Design: Optimized for float32 precision

Error Checking: Comprehensive tensor validation

Edge Case Handling: Empty tensor and size validation

PyTorch Integration: Seamless tensor passing and gradient support

Key CUDA Concepts Used
Shared Memory Reduction for parallel statistics computation

Grid-Stride Loops for workload distribution

Memory Coalescing for efficient global memory access

Numerical Stability through careful floating-point handling

Kernel Fusion combining multiple operations in single kernels

Expected Performance Benefits
2-5x speedup over PyTorch implementation for large batch sizes

Better numerical stability with log-sum-exp trick

Scalable performance with increasing batch dimensions

GPU utilization through optimized block and grid sizing

This implementation provides a production-ready, numerically stable Circle Loss with significant performance improvements through careful CUDA optimization and parallel reduction patterns.
Here's an example to show you the syntax of inline embedding custom CUDA operators in torch: The example given architecture is:
import torch
import torch.nn as nn
import torch.nn.functional as F

# 定义常量
BATCH_SIZE = 256
FEATURE_DIM = 512
MARGIN = 0.25
GAMMA = 256

class Model(nn.Module):

    def __init__(self):
        super().__init__()
        self.margin = MARGIN
        self.gamma = GAMMA

    def forward(self, features: torch.Tensor, labels: torch.Tensor) -> torch.Tensor:

        # Circle Loss的 PyTorch 实现
        # 假设 features 已经是 L2 归一化的
        # (B, D) @ (D, B) -> (B, B)
        similarities = torch.matmul(features, features.t())

        # 创建正负样本对的掩码
        mask_positive = labels.unsqueeze(1) == labels.unsqueeze(0)
        mask_negative = labels.unsqueeze(1) != labels.unsqueeze(0)

        # 收集正样本对和负样本对的相似度
        # .masked_select() 会将张量展平
        sp = similarities[mask_positive]
        sn = similarities[mask_negative]

        # 计算 Circle Loss 的 logits
        # .detach() 用于停止梯度反向传播
        ap = torch.clamp_min(-sp.detach() + 1 + self.margin, min=0.)
        an = torch.clamp_min(sn.detach() + self.margin, min=0.)

        delta_p = 1 - self.margin
        delta_n = self.margin

        logit_p = -ap * (sp - delta_p) * self.gamma
        logit_n = an * (sn - delta_n) * self.gamma

        # 使用 logsumexp 和 softplus 计算最终的 loss
        # 这是 "unified" 版本的 loss
        loss = F.softplus(torch.logsumexp(logit_n, dim=0) + torch.logsumexp(logit_p, dim=0))

        return loss

def get_inputs():

    # 特征需要 L2 归一化
    features = F.normalize(torch.randn(BATCH_SIZE, FEATURE_DIM, dtype=torch.float32), p=2, dim=1)
    labels = torch.randint(0, 10, (BATCH_SIZE,), dtype=torch.long) # 假设有 10 个类别
    return [features, labels]

def get_init_inputs():
    return []